package com.sh.data.engine.iotdb.client;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.EnumUtil;
import com.google.common.collect.Lists;
import com.sh.data.engine.iotdb.domain.DeviceDomain;
import com.sh.data.engine.iotdb.domain.TabletDomain;
import com.sh.data.engine.iotdb.domain.TabletDomain.TabletDataSet;
import com.sh.data.engine.iotdb.domain.TimeseriesDataSetDomain;
import com.sh.data.engine.iotdb.domain.TimeseriesDomain;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.iotdb.isession.SessionDataSet;
import org.apache.iotdb.isession.util.Version;
import org.apache.iotdb.rpc.IoTDBConnectionException;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.session.Session;
import org.apache.iotdb.tsfile.file.metadata.enums.CompressionType;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.apache.iotdb.tsfile.file.metadata.enums.TSEncoding;
import org.apache.iotdb.tsfile.read.common.RowRecord;
import org.apache.iotdb.tsfile.write.record.Tablet;
import org.apache.iotdb.tsfile.write.schema.MeasurementSchema;

import java.util.*;
import java.util.stream.Collectors;

@Slf4j
public class IoTDBSessionUtil {

    public static final String TIMESTAMP_MEASUREMENT = "Time(IOTDB)";

    public static final String TIMESTAMP_MEASUREMENT_2 = "time";

    public static Session getIoTDBSession(String server, String username, String password) {
        List<String> nodes = Arrays.asList(server.split(","));
        nodes = nodes.stream().filter(StringUtils::isNotBlank).collect(Collectors.toList());
        Session session = null;
        if (1 == nodes.size()) {
            String[] serverPort = nodes.get(0).split(":");
            session =
                new Session.Builder()
                    .host(serverPort[0])
                    .port(Integer.parseInt(serverPort[1]))
                    .username(username)
                    .password(password)
                    .version(Version.V_0_13)
                    .build();
        } else {
            session = new Session(nodes, username, password);
        }
        return session;
    }

    public static void main(String[] args)
        throws IoTDBConnectionException, StatementExecutionException {

        Session ioTDBSession = getIoTDBSession("127.0.0.1:6667", "root", "root");
//    ioTDBSession =
//            new Session.Builder()
//                    .host(serverPort[0])
//                    .port(Integer.parseInt(serverPort[1]))
//                    .username(username)
//                    .password(password)
//                    .version(Version.V_0_13)
//                    .build();
        ioTDBSession.open();
        String deviceId = "root.sestest.wf01.wt01";
        List<String> measurements = new ArrayList<>(3);
        measurements.add("s0");
        measurements.add("s1");
        measurements.add("s2");


        for (long i = 5000; i < 6000; i++) {
            List<String> values = new ArrayList<>(3);
            values.add("5");
            values.add("6");
            values.add("7");

            ioTDBSession.insertRecord(deviceId, i, measurements, values);
        }
        ioTDBSession.close();
    }


    /**
     * 获取时序
     *
     * @param pathPattern
     * @return
     * @throws Exception
     */
    public List<TimeseriesDomain> showTimeseries(Session session, String pathPattern)
        throws Exception {
        List<TimeseriesDomain> timeseriesList = Lists.newArrayList();
        // session.open(false, this.ioTDBOptions.getTimeout() * 1000);
        String sql = String.join(" ", "show timeseries", pathPattern);
        SessionDataSet dataSet = session.executeQueryStatement(sql);
        while (dataSet.hasNext()) {
            RowRecord next = dataSet.next();
            TimeseriesDomain timeseries = new TimeseriesDomain();
            timeseries.setTimeseries(next.getFields().get(0).getStringValue());
            timeseries.setAlias(next.getFields().get(1).getStringValue());
            timeseries.setStorageGroup(next.getFields().get(2).getStringValue());
            timeseries.setDataType(
                EnumUtil.fromString(TSDataType.class, next.getFields().get(3).getStringValue()));
            timeseries.setEncoding(
                EnumUtil.fromString(TSEncoding.class, next.getFields().get(4).getStringValue()));
            timeseries.setCompression(
                EnumUtil.fromString(CompressionType.class, next.getFields().get(5).getStringValue()));
            timeseries.setTags(next.getFields().get(6).getStringValue());
            timeseries.setAttributes(next.getFields().get(7).getStringValue());
            timeseriesList.add(timeseries);
        }
        return timeseriesList;
    }

    // 获取实体
    public static List<DeviceDomain> showDevices(Session session, String pathPattern)
        throws Exception {
        List<DeviceDomain> devices = Lists.newArrayList();
        pathPattern = Optional.ofNullable(pathPattern).orElse("");
        String sql = String.join(" ", "show devices", pathPattern);
        SessionDataSet dataSet = session.executeQueryStatement(sql);
        while (dataSet.hasNext()) {
            RowRecord next = dataSet.next();
            DeviceDomain deviceDomain = new DeviceDomain();
            deviceDomain.setDevice(next.getFields().get(0).getStringValue());
            String alignedValue = next.getFields().get(1).getStringValue();
            boolean isAligned = "true".equals(alignedValue) ? true : false;
            deviceDomain.setAligned(isAligned);
            devices.add(deviceDomain);
        }
        return devices;
    }

    // 插入时序数据-单条
    public static void insertRecordOneByOne(Session session, List<TimeseriesDataSetDomain> dataSets)
        throws IoTDBConnectionException, StatementExecutionException {
        for (TimeseriesDataSetDomain dataSet : dataSets) {
            String deviceId = dataSet.getDeviceId();
            long timestamp =
                Optional.ofNullable(dataSet.getTimestamp()).orElse(System.currentTimeMillis());
            session.insertRecord(
                deviceId,
                timestamp,
                dataSet.getMeasurements(),
                dataSet.getDataTypes(),
                dataSet.getValues());
        }
    }

    // 插入时序数据-单条 字符串格式
    public void insertStrRecordOneByOne(Session session, List<TimeseriesDataSetDomain> dataSets)
        throws IoTDBConnectionException, StatementExecutionException {
        for (TimeseriesDataSetDomain dataSet : dataSets) {
            String deviceId = dataSet.getDeviceId();
            long timestamp =
                Optional.ofNullable(dataSet.getTimestamp()).orElse(System.currentTimeMillis());
            List<String> values = Convert.toList(String.class, dataSet.getValues());
            session.insertRecord(deviceId, timestamp, dataSet.getMeasurements(), values);
        }
    }

    /*
     * A Tablet example:
     * device1
     * time s1, s2, s3
     * 1, 1, 1, 1
     * 2, 2, 2, 2
     * 3, 3, 3, 3
     */
    // 批量插入 指定批大小
    public static void insertTablet(Session session, TabletDomain tabletDomain) throws Exception {
        List<MeasurementSchema> schemaList = tabletDomain.getSchemaList();
        Tablet tablet =
            new Tablet(tabletDomain.getDeviceId(), schemaList, tabletDomain.getMaxRowNumber());
        boolean isAligned = false;
        List<DeviceDomain> devices = IoTDBSessionUtil.showDevices(session, tabletDomain.getDeviceId());
        if (CollectionUtils.isNotEmpty(devices)) {
            isAligned = devices.get(0).isAligned();
        }

        for (TabletDataSet tabletDataSet : tabletDomain.getTabletDataSets()) {
            int rowIndex = tablet.rowSize++;
            long timestamp =
                Optional.ofNullable(tabletDataSet.getTimestamp()).orElse(System.currentTimeMillis());
            tablet.addTimestamp(rowIndex, timestamp);
            List<Object> values = tabletDataSet.getValues();
            for (int i = 0; i < schemaList.size(); i++) {

                Object val = values.get(i);
                TSDataType type = schemaList.get(i).getType();
                switch (type) {
                    case BOOLEAN:
                        val = Convert.convert(Boolean.class, val);
                        break;
                    case INT32:
                        val = Convert.convert(Integer.class, val);
                        break;
                    case INT64:
                        val = Convert.convert(Long.class, val);
                        break;
                    case FLOAT:
                        val = Convert.convert(Float.class, val);
                        break;
                    case DOUBLE:
                        val = Convert.convert(Double.class, val);
                        break;
                    case TEXT:
                        val = Convert.convert(String.class, val);
                        break;
                    case VECTOR:
                        break;
                }
                tablet.addValue(schemaList.get(i).getMeasurementId(), rowIndex, val);
            }
            if (tablet.rowSize == tablet.getMaxRowNumber()) {
                if (isAligned) {
                    session.insertAlignedTablet(tablet, true);
                } else {
                    session.insertTablet(tablet, true);
                }
                tablet.reset();
            }
        }
        if (tablet.rowSize != 0) {
            if (isAligned) {
                session.insertAlignedTablet(tablet, true);
            } else {
                session.insertTablet(tablet, true);
            }
            tablet.reset();
        }
    }

    /**
     * @description : 算盘那里插入数据。测点数据集合values 里面的Object如果是JsonArray，那么要做一个拆分。
     * <p>比如同一时刻某个字段的值是 a,b,c 那么要插入3行而不是一行
     * @author : tinglan.ys
     * @date : 2022/12/9
     */

    // TabletDomain(deviceId=root.sg1.d3,
    // schemaList=[[s3,TEXT,PLAIN,,SNAPPY],
    // [s1,TEXT,PLAIN,,SNAPPY],
    // [s2,TEXT,PLAIN,,SNAPPY]],
    // maxRowNumber=100,
    // tabletDataSets=[
    // TabletDomain.TabletDataSet(values=
    // [ys,
    // ["reference","fiction","fiction","fiction"],
    // ["Nigel Rees","Evelyn Waugh","Herman Melville","J. R. R. Tolkien"]],
    // timestamp=1670581899996)])
    // 其中timestamp为这一批数据入iotdb的时间，这个一般都用系统时间就行了，如果要存时间，那么再开一个另外的字段去存时间
    public static void insertTabletSp(Session session, TabletDomain tabletDomain) throws Exception {
        List<MeasurementSchema> schemaList = tabletDomain.getSchemaList();
        Tablet tablet = new Tablet(tabletDomain.getDeviceId(), schemaList);

        boolean isAligned = false;
        List<DeviceDomain> devices = IoTDBSessionUtil.showDevices(session, tabletDomain.getDeviceId());
        if (CollectionUtils.isNotEmpty(devices)) {
            isAligned = devices.get(0).isAligned();
        }

        TabletDataSet tabletDataSet = tabletDomain.getTabletDataSets().get(0);

        // 拿到的demo数据见注释
        // timestamp是固定的 因为是同一时刻进来的数据。后来发现这个timestamp不能传一样的值，需要做个timestamp++，坑！
        long timestamp =
            Optional.ofNullable(tabletDataSet.getTimestamp()).orElse(System.currentTimeMillis());

        List<Object> values = tabletDataSet.getValues();

        // allListRowNum为每个字段的行数
        List<Integer> allListRowNum = new ArrayList<>();

        for (int i = 0; i < schemaList.size(); i++) {
            // 每一个字段rowIndex都从0开始
            int rowIndex = 0;
            Object val = values.get(i);
            Object insertVal = null;
            List insertValList = Convert.convert(List.class, val);
            allListRowNum.add(insertValList.size());
            for (Object o : insertValList) {
                TSDataType type = schemaList.get(i).getType();
                switch (type) {
                    case BOOLEAN:
                        insertVal = Convert.convert(Boolean.class, o);
                        break;
                    case INT32:
                        insertVal = Convert.convert(Integer.class, o);
                        break;
                    case INT64:
                        insertVal = Convert.convert(Long.class, o);
                        break;
                    case FLOAT:
                        insertVal = Convert.convert(Float.class, o);
                        break;
                    case DOUBLE:
                        insertVal = Convert.convert(Double.class, o);
                        break;
                    case TEXT:
                        insertVal = Convert.convert(String.class, o);
                        break;
                    case VECTOR:
                        break;
                    default:
                }
                tablet.addValue(schemaList.get(i).getMeasurementId(), rowIndex, insertVal);
                rowIndex++;
            }
        }

        Integer max = Collections.max(allListRowNum);
        // 设置行数
        tablet.rowSize = Collections.max(allListRowNum);

        if (tablet.rowSize != 0) {
            if (isAligned) {
                for (int i = 0; i < max; i++) {
                    // 设置时间戳，注意这里不能是一样的！！！
                    tablet.addTimestamp(i, timestamp);
                }
                session.insertAlignedTablet(tablet, true);
            } else {
                for (int i = 0; i < max; i++) {
                    // 设置时间戳，注意这里不能是一样的！！！
                    tablet.addTimestamp(i, timestamp);
                    timestamp++;
                }
                session.insertTablet(tablet, true);
            }
            tablet.reset();
        }
    }

    public static void insertTabletSpV2(
        Session session, TabletDomain tabletDomain, Boolean useSysTimeFlag, List<Long> timestamps)
        throws Exception {

        List<MeasurementSchema> schemaList = tabletDomain.getSchemaList();
        Tablet tablet = new Tablet(tabletDomain.getDeviceId(), schemaList);

        boolean isAligned = false;
        List<DeviceDomain> devices = IoTDBSessionUtil.showDevices(session, tabletDomain.getDeviceId());
        if (CollectionUtils.isNotEmpty(devices)) {
            isAligned = devices.get(0).isAligned();
        }

        List<Object> values = tabletDomain.getInsertValues();

        // allListRowNum为每个字段的行数
        List<Integer> allListRowNum = new ArrayList<>();

        for (int i = 0; i < schemaList.size(); i++) {
            // 每一个字段rowIndex都从0开始
            int rowIndex = 0;
            Object val = values.get(i);
            Object insertVal = null;
            List insertValList = Convert.convert(List.class, val);
            allListRowNum.add(insertValList.size());
            for (Object o : insertValList) {
                TSDataType type = schemaList.get(i).getType();
                switch (type) {
                    case BOOLEAN:
                        insertVal = Convert.convert(Boolean.class, o);
                        break;
                    case INT32:
                        insertVal = Convert.convert(Integer.class, o);
                        break;
                    case INT64:
                        insertVal = Convert.convert(Long.class, o);
                        break;
                    case FLOAT:
                        insertVal = Convert.convert(Float.class, o);
                        break;
                    case DOUBLE:
                        insertVal = Convert.convert(Double.class, o);
                        break;
                    case TEXT:
                        insertVal = Convert.convert(String.class, o);
                        break;
                    case VECTOR:
                        break;
                    default:
                }
                tablet.addValue(schemaList.get(i).getMeasurementId(), rowIndex, insertVal);
                rowIndex++;
            }
        }

        Integer max = Collections.max(allListRowNum);
        // 设置行数
        tablet.rowSize = Collections.max(allListRowNum);

        if (useSysTimeFlag) {
            // 使用系统时间
            long timestamp = System.currentTimeMillis();
            for (int i = 0; i < max; i++) {
                tablet.addTimestamp(i, timestamp);
                timestamp++;
            }
        } else {
            for (int i = 0; i < timestamps.size(); i++) {
                // tips:注意如果timestamps的值是相同的，那么插入的时候会出错。强烈建议使用系统时间
                tablet.addTimestamp(i, timestamps.get(i));
            }
        }

        if (tablet.rowSize != 0) {
            if (isAligned) {
                session.insertAlignedTablet(tablet, true);
            } else {
                session.insertTablet(tablet, true);
            }
            tablet.reset();
        }
    }
}
